package com.sheepone.util.weka;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import weka.core.*;
import weka.core.converters.ArffSaver;
import weka.core.converters.ConverterUtils.DataSource;

import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;

/**
 * @author Miss.杨
 * @description
 * @since 2024/3/4 - 21:36
 */

public class WekaUtil {
    private static final Logger logger = LoggerFactory.getLogger(new Object() {
    }.getClass().getEnclosingClass());

    /**
     * 将double[][]数据转换为Weka Instances对象
     *
     * @param data
     * @return
     * @throws Exception
     */
    public static Instances convertToInstances(double[][] data) {
        // No class attribute for clustering
        // -1 表明这是无监督学习任务
        return convertToInstances(data, -1);
    }

    /**
     * 将double[][]数据转换为Weka Instances对象
     *
     * @param data
     * @param classIndex
     * @return
     * @throws Exception
     */
    public static Instances convertToInstances(double[][] data, int classIndex) {
        ArrayList<Attribute> attributes = new ArrayList<>();
        for (int i = 0; i < data[0].length; i++) {
            attributes.add(new Attribute("Attribute" + i));
        }
        Instances instances = new Instances("Dataset", attributes, data.length);
        for (double[] instance : data) {
            Instance wekaInstance = convertToInstance(instance);
            instances.add(wekaInstance);
        }
        instances.setClassIndex(classIndex);
        return instances;
    }

    /**
     * 将Double[][]数据转换为Weka Instances对象
     *
     * @param data
     * @param classIndex
     * @return
     * @throws Exception
     */
    public static Instances convertToInstances(Double[][] data, int classIndex) {
        ArrayList<Attribute> attributes = new ArrayList<>();
        for (int i = 0; i < data[0].length; i++) {
            attributes.add(new Attribute("Attribute" + i));
        }
        Instances instances = new Instances("Dataset", attributes, data.length);
        for (Double[] instance : data) {
            Instance wekaInstance = convertToInstance(instance);
            instances.add(wekaInstance);
        }
        instances.setClassIndex(classIndex);
        return instances;
    }

    public static Instance convertToInstance(double[] instance) {
        Instance wekaInstance = new DenseInstance(instance.length);
        for (int j = 0; j < instance.length; j++) {
            wekaInstance.setValue(j, instance[j]);
        }
        return wekaInstance;
    }

    public static Instance convertToInstance(Double[] instance) {
        Instance wekaInstance = new DenseInstance(instance.length);
        for (int j = 0; j < instance.length; j++) {
            wekaInstance.setValue(j, instance[j]);
        }
        return wekaInstance;
    }

    /**
     * 计算两个实例之间的欧氏距离
     *
     * @param instance1 第一个实例
     * @param instance2 第二个实例
     * @return 两实例间的欧氏距离
     */
    public static double calculateEuclideanDistance(Instance instance1, Instance instance2) {
        double sum = 0.0;
        for (int i = 0; i < instance1.numAttributes(); i++) {
            if (!instance1.isMissing(i) && !instance2.isMissing(i)) {
                double diff = instance1.value(i) - instance2.value(i);
                sum += Math.pow(diff, 2);
            }
        }
        return Math.sqrt(sum);
    }

    public static boolean saveInstances(Instances instances, String filePath) {
        if (!filePath.endsWith(".arff")) {
            filePath = filePath + ".arff";
        }
        try {
            ArffSaver saver = new ArffSaver();
            saver.setInstances(instances);
            saver.setFile(new File(filePath));
            saver.writeBatch();
            return true;
        } catch (Exception e) {
            logger.error("save Instances to {} fail", filePath);
            return false;
        }
    }

    public static Instances loadInstancesFromArff(String filePath) {
        // 使用Paths处理文件路径，确保路径处理的正确性
        Path path = Paths.get(filePath);
        String fileName = path.getFileName().toString();
        if (!fileName.endsWith(".arff")) {
            filePath = path.getParent().toString() + File.separator + fileName + ".arff";
        }
        File file = new File(filePath);
        // 检查文件是否存在
        path = file.toPath();
        if (!Files.exists(path)) {
            logger.error("load Instances from {} fail", filePath);
            return null;
        }
        try {
            // 创建DataSource对象，用于读取.arff文件
            DataSource dataSource = new DataSource(filePath);
            // 加载.arff文件内容并转换为Instances对象
            return dataSource.getDataSet();
        } catch (Exception e) {
            logger.error("load Instances from {} fail", filePath);
            return null;
        }
    }

}
